import vtk
from config import Scale_Percent_Setting
from components.dicom_tag_info import DicomTagInfo
from utils.util import img_to_png_base64, img_to_jpeg_base64
import numpy as np


def set_origin(dicom_reader, actor):
    bounds = dicom_reader.GetOutput().GetBounds()
    center = [(bounds[0] + bounds[1]) / 2, (bounds[2] + bounds[3]) / 2, (bounds[4] + bounds[5]) / 2]
    actor.SetOrigin(center)


def create_tissue_presets():
    presets = {
        'bone': {
            'color_points': [
                (150, 1.0, 1.0, 1.0),  # 白色
                (400, 1.0, 1.0, 1.0),  # 白色
                (650, 1.0, 1.0, 1.0),  # 白色
            ],
            'opacity_points': [
                (150, 0.0),
                (400, 0.5),
                (650, 1.0),
            ]
        },
        'AAA': {
            'color_points': [
                (-3024, 0.0, 0.0, 0.0),  # 黑色
                (150, 0.2, 0.2, 0.2),  # 深灰色
                (400, 0.8, 0.8, 0.8),  # 浅灰色
                (650, 1.0, 1.0, 1.0),  # 白色
            ],
            'opacity_points': [
                # (-100, 0.0),
                # (100, 0.1),  # 低密度区域几乎透明
                # (300, 0.3),  # 血管区域半透明
                # (500, 0.4),
                (-3024, 0.0),
                (150, 0.0),  # 低密度区域完全透明
                (300, 0.5),  # 中等密度区域稍微显示
                (500, 1.0),  # 高密度区域部分显示
            ]
        },
        'blood': {
            'color_points': [
                (120, 0.2, 0.0, 0.0),
                (170, 0.4, 0.0, 0.0),
                (220, 0.6, 0.2, 0.0),
                (270, 0.7, 0.2, 0.0),
                (370, 0.8, 0.2, 0.0),
            ],
            'opacity_points': [
                (120, 0.0),
                (170, 0.2),
                (220, 0.3),
                (270, 0.2),
                (370, 0.1),
            ]
        },
    }
    return presets


def apply_preset(color_tf, opacity_tf, preset):
    # 清除现有的点
    color_tf.RemoveAllPoints()
    opacity_tf.RemoveAllPoints()

    # 添加颜色点
    for point in preset['color_points']:
        color_tf.AddRGBPoint(point[0], point[1], point[2], point[3])

    # 添加不透明度点
    for point in preset['opacity_points']:
        opacity_tf.AddPoint(point[0], point[1])


def setup_lighting(renderer):
    # 添加环境光
    ambient_light = vtk.vtkLight()
    ambient_light.SetLightTypeToHeadlight()
    ambient_light.SetIntensity(1)
    ambient_light.SetColor(1, 1, 1)
    renderer.AddLight(ambient_light)

    # 添加漫反射光
    diffuse_light = vtk.vtkLight()
    diffuse_light.SetLightTypeToSceneLight()
    diffuse_light.SetPosition(0, 1, 1)
    diffuse_light.SetIntensity(0.8)
    diffuse_light.SetColor(1, 1, 1)
    renderer.AddLight(diffuse_light)


class VolumeLoader:
    def __init__(self, dicom_tag_info):
        self.render_window = None
        self.windowToImageFilter = None
        self.writer = None
        self.dicom_folder_path = ''
        self.dicom_tag_info: DicomTagInfo = dicom_tag_info
        self.main_viewport = {
            'renderer': None,
            "volume_mapper": None,
            "color_transfer_function": None,
            "opacity_transfer_function": None,
            "volume_property": None,
            "transform": vtk.vtkTransform(),
            "actor": None,
            "template": "",
            "window": [],
        }
        self.minor_viewport = {
            "renderer": None,
            "transform": vtk.vtkTransform(),
            "actor": None
        }

        self.start_boost_flag = False
        self.create_obj()
        self.init_main_render()
        self.init_minor_render()
        self.init_write()
        self.init_render_window()
        self.start_boost()
        self.init_cube()

    def create_obj(self):
        self.render_window = vtk.vtkRenderWindow()
        self.main_viewport['volume_mapper'] = vtk.vtkGPUVolumeRayCastMapper()
        self.main_viewport["volume_property"] = vtk.vtkVolumeProperty()
        self.main_viewport["color_transfer_function"] = vtk.vtkColorTransferFunction()
        self.main_viewport["opacity_transfer_function"] = vtk.vtkPiecewiseFunction()

    def init_main_render(self):
        renderer = vtk.vtkRenderer()
        renderer.SetViewport(0.0, 0.0, 1.0, 1.0)
        self.main_viewport['renderer'] = renderer

    def init_minor_render(self):
        renderer = vtk.vtkRenderer()
        renderer.SetViewport(0.9, 0.9, 1.0, 1.0)
        self.minor_viewport['renderer'] = renderer

    def init_render_window(self):
        self.render_window.SetOffScreenRendering(1)
        self.render_window.SetDesiredUpdateRate(60.0)
        self.render_window.AddRenderer(self.main_viewport['renderer'])
        self.render_window.AddRenderer(self.minor_viewport['renderer'])

    def set_render_size(self, size):
        self.render_window.SetSize(size)

    def start_boost(self):
        # 开启抗锯齿
        self.render_window.SetMultiSamples(1)

        main_volume_mapper = self.main_viewport['volume_mapper']
        main_volume_mapper.SetUseJittering(True)
        main_volume_mapper.SetSampleDistance(3.0)  # 增加采样距离可提高性能
        # 设置图像采样率
        main_volume_mapper.SetImageSampleDistance(3.0)
        # 自动调整采样率
        main_volume_mapper.SetAutoAdjustSampleDistances(False)
        # 强制更新
        self.writer.Modified()  # 标记写入器已修改
        main_volume_mapper.Modified()
        self.start_boost_flag = True

    def stop_boost(self):
        self.render_window.SetMultiSamples(1)
        main_volume_mapper = self.main_viewport['volume_mapper']
        main_volume_mapper.SetSampleDistance(1.0)  # 增加采样距离可提高性能
        # 设置图像采样率
        main_volume_mapper.SetImageSampleDistance(1.0)
        # 自动调整采样率
        main_volume_mapper.SetAutoAdjustSampleDistances(False)
        # 强制更新
        self.writer.Modified()  # 标记写入器已修改
        main_volume_mapper.Modified()
        self.start_boost_flag = False

    def apply_base_window(self, color_points, opacity_points):
        wl = int(self.dicom_tag_info.tags['window_center'])
        ww = int(self.dicom_tag_info.tags['window_width'])
        color_points.append((wl - ww / 2, 0.0, 0.0, 0.0))
        color_points.append((wl, 1.0, 1.0, 1.0))
        color_points.append((wl + ww / 2, 0.5, 0.5, 0.5))

        opacity_points.append((wl - ww / 2, 0.0))
        opacity_points.append((wl, 0.5))
        opacity_points.append((wl + ww / 2, 0.0))

    def set_transfer_function(self):
        color_transfer_function = self.main_viewport['color_transfer_function']
        opacity_transfer_function = self.main_viewport['opacity_transfer_function']
        color_transfer_function.RemoveAllPoints()
        opacity_transfer_function.RemoveAllPoints()

        volume_property = self.main_viewport["volume_property"]
        presets = create_tissue_presets()
        template_name = self.main_viewport['template']
        window_arr = self.main_viewport['window']
        color_points = []
        opacity_points = []
        if template_name == '' and len(window_arr) == 0:
            self.apply_base_window(color_points, opacity_points)
        elif template_name != '':
            color_points = presets[template_name]['color_points']
            opacity_points = presets[template_name]['opacity_points']
        for window in window_arr:
            color_points = color_points + presets[window]['color_points']
            opacity_points = opacity_points + presets[window]['opacity_points']
        combined_preset = {
            'color_points': color_points,
            'opacity_points': opacity_points
        }
        apply_preset(color_transfer_function, opacity_transfer_function, combined_preset)
        volume_property.SetColor(color_transfer_function)
        volume_property.SetScalarOpacity(opacity_transfer_function)
        volume_property.Modified()

    def load_dicom_series(self, dicom_folder_path):
        self.dicom_folder_path = dicom_folder_path
        dicom_reader = vtk.vtkDICOMImageReader()
        dicom_reader.SetDirectoryName(self.dicom_folder_path)
        dicom_reader.Update()
        volume_mapper = self.main_viewport['volume_mapper']
        volume_property = self.main_viewport["volume_property"]
        volume_mapper.SetInputConnection(dicom_reader.GetOutputPort())
        volume_property.SetInterpolationTypeToLinear()
        self.set_transfer_function()
        volume_property.ShadeOn()  # 开启阴影
        volume_property.SetAmbient(0.4)
        volume_property.SetDiffuse(0.6)
        volume_property.SetSpecular(0.3)  # 增加高光
        volume_property.SetSpecularPower(5)  # 调整高光强度
        # 5. 创建体积对象并设置映射器和属性
        actor = vtk.vtkVolume()
        actor.SetMapper(volume_mapper)
        actor.SetProperty(volume_property)
        self.main_viewport['actor'] = actor
        main_renderer = self.main_viewport['renderer']
        main_renderer.AddVolume(actor)
        main_renderer.SetBackground(0.0, 0, 0)
        setup_lighting(main_renderer)
        set_origin(dicom_reader, actor)

    def get_recent_img(self):

        base64_img = self.get_base64()
        return {
            "img": base64_img
        }

    def get_high_quality_img(self):
        self.stop_boost()
        self.main_viewport['volume_mapper'].Update()
        self.render_window.Render()
        clear_img = self.get_recent_img()
        self.start_boost()
        self.main_viewport['volume_mapper'].Update()
        self.render_window.Render()
        return clear_img

    def init_write(self):
        self.writer = vtk.vtkPNGWriter()
        self.writer.SetWriteToMemory(1)
        self.windowToImageFilter = vtk.vtkWindowToImageFilter()
        self.windowToImageFilter.SetInputBufferTypeToRGB()
        self.windowToImageFilter.SetInput(self.render_window)

    def get_base64(self):
        self.render_window.Render()
        self.windowToImageFilter.Modified()  # 标记过滤器已修改
        self.writer.SetInputConnection(self.windowToImageFilter.GetOutputPort())
        self.writer.Modified()  # 标记写入器已修改
        self.writer.Write()
        img_bytes = bytes(memoryview(self.writer.GetResult()))
        img_base64 = img_to_png_base64(img_bytes)
        return img_base64

    def update_pan(self, delta_x, delta_y):
        self.main_viewport['transform'].Translate(delta_x, -delta_y, 0)
        self.main_viewport['actor'].SetUserTransform(self.main_viewport['transform'])

    def update_zoom(self, zoom_factor):
        main_transform = self.main_viewport['transform']
        main_actor = self.main_viewport['actor']
        current_scale = main_transform.GetScale()
        target_scale = np.clip(zoom_factor + current_scale[0], Scale_Percent_Setting[0], Scale_Percent_Setting[1])
        target_scale = target_scale / current_scale[0]
        main_transform.Scale(target_scale, target_scale, target_scale)
        main_actor.SetUserTransform(main_transform)

    def update_rotate(self, x, y):
        main_actor = self.main_viewport['actor']
        minor_actor = self.minor_viewport['actor']
        if x != 0:
            main_actor.RotateY(x)
            minor_actor.RotateY(x)
        if y != 0:
            main_actor.RotateX(y)
            minor_actor.RotateX(y)

    def init_cube(self):
        cube = vtk.vtkCubeSource()
        # 创建一个Assembly作为父节点
        self.minor_viewport['actor'] = vtk.vtkAssembly()
        minor_actor = self.minor_viewport['actor']
        minor_renderer = self.minor_viewport['renderer']

        cube.SetXLength(1)
        cube.SetYLength(1)
        cube.SetZLength(1)
        cube.Update()
        # 定义每个面的颜色
        face_colors = [
            (1, 0, 0),  # 红色 - 前面
            (0, 1, 0),  # 绿色 - 后面
            (0, 0, 1),  # 蓝色 - 左面
            (1, 1, 0),  # 黄色 - 右面
            (1, 0, 1),  # 品红 - 上面
            (0, 1, 1)  # 青色 - 下面
        ]

        # 获取立方体的 PolyData
        poly_data = cube.GetOutput()

        # 为每个面指定颜色
        colors = vtk.vtkUnsignedCharArray()
        colors.SetNumberOfComponents(3)  # RGB
        colors.SetName("Colors")

        # 立方体有6个面，每个面添加一个颜色
        for color in face_colors:
            colors.InsertNextTuple3(*(int(c * 255) for c in color))

        # 将颜色赋给立方体的 CellData
        poly_data.GetCellData().SetScalars(colors)

        labels = ['I', 'S', 'R', 'L', 'A', 'P']
        text_actors = []

        for i, label in enumerate(labels):
            text = vtk.vtkVectorText()
            text.SetText(label)

            # 文本映射器和actor
            text_mapper = vtk.vtkPolyDataMapper()
            text_mapper.SetInputConnection(text.GetOutputPort())

            text_actor = vtk.vtkActor()
            text_actor.SetMapper(text_mapper)
            text_actor.GetProperty().SetColor(255, 255, 255)
            text_actor.SetScale(0.7)
            delta_x = -0.45
            delta_y = -0.35
            scale = 0.51
            # 定位文字
            transform = vtk.vtkTransform()
            if i == 0:  # 前面
                transform.Translate(delta_x, delta_y, scale)
            elif i == 1:  # 后面
                transform.RotateY(180)
                transform.Translate(delta_x, delta_y, scale)
            elif i == 2:  # 左面
                transform.RotateY(-90)
                transform.Translate(delta_x, delta_y, scale)
            elif i == 3:  # 右面
                transform.RotateY(90)
                transform.Translate(delta_x, delta_y, scale)
            elif i == 4:  # 上面
                transform.RotateX(-90)
                transform.Translate(delta_x, delta_y, scale)
            elif i == 5:  # 下面
                transform.RotateX(90)
                transform.Translate(delta_x, delta_y, scale)

            text_actor.SetUserTransform(transform)
            text_actors.append(text_actor)

        for text_actor in text_actors:
            minor_actor.AddPart(text_actor)

        mapper = vtk.vtkPolyDataMapper()
        mapper.SetInputData(poly_data)
        cube_actor = vtk.vtkActor()
        cube_actor.SetMapper(mapper)

        minor_actor.AddPart(cube_actor)
        minor_renderer.AddActor(minor_actor)

    def set_3d_params(self, template, windows):
        self.main_viewport["template"] = template
        self.main_viewport['window'] = windows
        self.set_transfer_function()

    def set_3d_direction(self, direction):
        main_camera = self.main_viewport['renderer'].GetActiveCamera()
        minor_camera = self.minor_viewport['renderer'].GetActiveCamera()
        main_actor = self.main_viewport['actor']
        minor_actor = self.minor_viewport['actor']
        main_actor.SetOrientation(0, 0, 0)
        minor_actor.SetOrientation(0, 0, 0)

        for camera in [main_camera, minor_camera]:
            if direction == 'A':  # Anterior
                camera.SetPosition(0, 1000, 0)
                camera.SetViewUp(0, 0, -1)
            elif direction == 'P':  # Posterior
                camera.SetPosition(0, -1000, 0)
                camera.SetViewUp(0, 0, -1)

            elif direction == 'S':  # Superior
                camera.SetPosition(0, 0, -1000)
                camera.SetViewUp(0, 1, 0)
            elif direction == 'I':  # Inferior
                camera.SetPosition(0, 0, 1000)
                camera.SetViewUp(0, 1, 0)

            elif direction == 'R':  # Right
                camera.SetPosition(-1000, 0, 0)
                camera.SetViewUp(0, 0, -1)
            elif direction == 'L':  # Left
                camera.SetPosition(1000, 0, 0)
                camera.SetViewUp(0, 0, -1)
            camera.SetFocalPoint(0, 0, 0)
        self.main_viewport['renderer'].ResetCamera()
        self.minor_viewport['renderer'].ResetCamera()

    def reset(self, reset_type):
        main_transform = self.main_viewport['transform']
        minor_transform = self.minor_viewport['transform']
        main_transform.Identity()  # 重置为单位矩阵
        minor_transform.Identity()
        self.main_viewport['actor'].SetUserTransform(main_transform)
        self.minor_viewport['actor'].SetUserTransform(minor_transform)
        # 重置Actor的位置和方向
        main_actor = self.main_viewport['actor']
        minor_actor = self.minor_viewport['actor']

        # 重置位置
        main_actor.SetPosition(0, 0, 0)
        minor_actor.SetPosition(0, 0, 0)

        # 重置方向
        main_actor.SetOrientation(0, 0, 0)
        minor_actor.SetOrientation(0, 0, 0)
